package cn.shengchao.examstar.ai.infrastructure.persistence.repository;

import cn.hutool.core.date.DateTime;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import cn.shengchao.examstar.ai.domain.repository.IAiSessionRepository;
import cn.shengchao.examstar.ai.infrastructure.dao.AiChatMemoryDao;
import cn.shengchao.examstar.ai.infrastructure.dao.AiSessionDao;
import cn.shengchao.examstar.ai.infrastructure.dao.po.AiChatMemory;
import cn.shengchao.examstar.ai.infrastructure.dao.po.AiSession;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Repository;
import top.continew.starter.cache.redisson.util.RedisUtils;

import java.time.Duration;
import java.util.Date;
import java.util.List;
import java.util.concurrent.CompletableFuture;


@Repository
public class AiSessionRepository implements IAiSessionRepository {
    @Autowired
    private AiSessionDao aiSessionDao;

    @Autowired
    private AiChatMemoryDao aiChatMemoryDao;


    @Override
    public AiSession createSession(String conversationId, Long userId, String type, Long objectId, Long subObjectId) {
        AiSession aiSession = new AiSession();
        if (StrUtil.isNotEmpty(conversationId)) {
            aiSession.setConversationId(conversationId);
        } else {
            aiSession.setConversationId(IdUtil.simpleUUID());
        }
        aiSession.setUserId(userId);
        aiSession.setType(type);
        aiSession.setObjectId(objectId);
        aiSession.setSubObjectId(subObjectId);
        aiSession.setCreateTime(new Date());
        aiSessionDao.insert(aiSession);
        return aiSession;
    }

    @Override
    public void saveChatHistory(String sessionId, String messages) {
        //写入缓存
        RedisUtils.set(RedisUtils.formatKey("chat_history", sessionId), messages, Duration.ofDays(1));
        //异步入库
        CompletableFuture.runAsync(() -> {
            //查询会话
            AiSession aiSession = aiSessionDao.selectOne(Wrappers.lambdaQuery(AiSession.class).eq(AiSession::getConversationId, sessionId));
            if (aiSession == null) {
                return;
            }
            //查询历史记录
            AiChatMemory aiChatMemory = aiChatMemoryDao.selectOne(Wrappers.lambdaQuery(AiChatMemory.class).eq(AiChatMemory::getConversationId, sessionId).last("limit 1"));
            if (aiChatMemory == null) {
                aiChatMemory = new AiChatMemory();
                aiChatMemory.setTimestamp(new Date());
            }
            aiChatMemory.setConversationId(sessionId);
            aiChatMemory.setMessages(messages);
            aiChatMemory.setUserId(aiSession.getUserId());
            aiChatMemoryDao.insertOrUpdate(aiChatMemory);
        });
    }

    @Override
    public String getChatHistory(String sessionId) {
        String messages = RedisUtils.get(RedisUtils.formatKey("chat_history", sessionId));
        if (messages == null) {
            AiChatMemory aiChatHistories = aiChatMemoryDao.selectOne(Wrappers.lambdaQuery(AiChatMemory.class)
                    .eq(AiChatMemory::getConversationId, sessionId));
            if (aiChatHistories != null) {
                RedisUtils.set(RedisUtils.formatKey("chat_history", sessionId), aiChatHistories.getMessages(), Duration.ofDays(1));
            }
        }
        return messages;
    }

    @Override
    public AiSession getSessionByConversationId(String conversationId) {
        return aiSessionDao.selectOne(Wrappers.lambdaQuery(AiSession.class)
                .eq(AiSession::getConversationId, conversationId));
    }

    @Override
    public void deleteSession(String conversationId) {
        // 删除redis缓存
        RedisUtils.set(RedisUtils.formatKey("chat_history", conversationId), null, Duration.ofDays(1));
        // 删除历史
        aiChatMemoryDao.delete(Wrappers.lambdaQuery(AiChatMemory.class).eq(AiChatMemory::getConversationId, conversationId));
    }

    @Override
    public Long getSessionByTime(Long userId, DateTime startTime, Date endTime) {
        return aiSessionDao.selectCount(Wrappers.lambdaQuery(AiSession.class)
                .ge(startTime != null, AiSession::getCreateTime, startTime)
                .le(endTime != null, AiSession::getCreateTime, endTime)
                .eq(AiSession::getUserId, userId));
    }

    @Override
    public List<String> getHistorySession(Long userId, Long objectId, Long subjectId) {
        if (objectId != null && subjectId != null) {
            List<AiSession> sessions = aiSessionDao.selectList(Wrappers.lambdaQuery(AiSession.class).eq(AiSession::getUserId, userId)
                    .eq(AiSession::getObjectId, objectId).eq(AiSession::getSubObjectId, subjectId));
            if (sessions.isEmpty()) {
                return List.of();
            }
            List<AiChatMemory> aiChatHistories = aiChatMemoryDao.selectList(Wrappers.lambdaQuery(AiChatMemory.class)
                    .in(AiChatMemory::getConversationId, sessions.stream().map(AiSession::getConversationId).toList())
                    .orderByAsc(AiChatMemory::getTimestamp));
            return aiChatHistories.stream().map(AiChatMemory::getMessages
            ).toList();
        }
        return List.of();
    }

    @Override
    public List<String> getHistorySessionByType(Long userId, String sessionType, Long objectId, Long subjectId) {
        if (objectId != null && subjectId != null) {
            List<AiSession> sessions = aiSessionDao.selectList(Wrappers.lambdaQuery(AiSession.class)
                    .eq(AiSession::getUserId, userId)
                    .eq(AiSession::getType, sessionType)
                    .eq(AiSession::getObjectId, objectId)
                    .eq(AiSession::getSubObjectId, subjectId));
            if (sessions.isEmpty()) {
                return List.of();
            }
            List<AiChatMemory> aiChatHistories = aiChatMemoryDao.selectList(Wrappers.lambdaQuery(AiChatMemory.class)
                    .in(AiChatMemory::getConversationId, sessions.stream().map(AiSession::getConversationId).toList())
                    .orderByAsc(AiChatMemory::getTimestamp));
            return aiChatHistories.stream().map(AiChatMemory::getMessages).toList();
        }
        return List.of();
    }

    @Override
    public List<String> getHistorySessionByTypePrefix(Long userId, String sessionTypePrefix, Long objectId, Long subjectId) {
        if (objectId != null && subjectId != null) {
            List<AiSession> sessions = aiSessionDao.selectList(Wrappers.lambdaQuery(AiSession.class)
                    .eq(AiSession::getUserId, userId)
                    .likeRight(AiSession::getType, sessionTypePrefix)
                    .eq(AiSession::getObjectId, objectId)
                    .eq(AiSession::getSubObjectId, subjectId));
            if (sessions.isEmpty()) {
                return List.of();
            }
            List<AiChatMemory> aiChatHistories = aiChatMemoryDao.selectList(Wrappers.lambdaQuery(AiChatMemory.class)
                    .in(AiChatMemory::getConversationId, sessions.stream().map(AiSession::getConversationId).toList())
                    .orderByAsc(AiChatMemory::getTimestamp));
            return aiChatHistories.stream().map(AiChatMemory::getMessages).toList();
        }
        return List.of();
    }

    @Override
    public boolean isSessionOwner(String conversationId, Long userId) {
        AiSession session = aiSessionDao.selectOne(Wrappers.lambdaQuery(AiSession.class)
                .eq(AiSession::getConversationId, conversationId)
                .eq(AiSession::getUserId, userId));
        return session != null;
    }
} 